package org.ruoyi.interview.controller;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.ruoyi.common.core.domain.model.LoginUser;
import org.ruoyi.common.satoken.utils.LoginHelper;
import org.ruoyi.interview.config.Ultra;
import org.ruoyi.interview.config.XunfeiSparkConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 面试AI WebSocket控制器
 * 提供前端直接通过WebSocket连接讯飞API的能力
 *
 * @author ruoyi
 */
@Slf4j
@RequiredArgsConstructor
@RestController
@RequestMapping("/interview/ai/websocket")
public class InterviewAIWebSocketController extends TextWebSocketHandler {
    
    private final XunfeiSparkConfig xunfeiSparkConfig;
    
    // 存储WebSocket会话和对应的讯飞客户端
    private final Map<String, Ultra.XunfeiWebSocketClientStream> sessionMap = new ConcurrentHashMap<>();
    
    /**
     * 获取WebSocket连接地址
     * 前端可以通过返回的URL建立WebSocket连接
     *
     * @return WebSocket连接URL
     */
    @GetMapping("/url")
    public String getWebSocketUrl() {
        try {
            Ultra ultra = new Ultra();
            ultra.setApiKey(xunfeiSparkConfig.getApiKey());
            ultra.setApiSecret(xunfeiSparkConfig.getApiSecret());
            ultra.setAppId(xunfeiSparkConfig.getAppId());
            
            String authUrl = ultra.getAuthUrl();
            return authUrl;
        } catch (Exception e) {
            log.error("获取WebSocket URL失败", e);
            return "error: " + e.getMessage();
        }
    }
    
    /**
     * 建立WebSocket连接时的处理
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("WebSocket连接已建立，会话ID: {}", session.getId());
        
        // 创建讯飞WebSocket客户端
        Ultra ultra = new Ultra();
        ultra.setApiKey(xunfeiSparkConfig.getApiKey());
        ultra.setApiSecret(xunfeiSparkConfig.getApiSecret());
        ultra.setAppId(xunfeiSparkConfig.getAppId());
        
        // 创建自定义的流式回调处理器
        Ultra.XunfeiWebSocketClientStream client = new Ultra.XunfeiWebSocketClientStream(
            xunfeiSparkConfig.getAppId(), 
            new AIWebSocketStreamCallback(session)
        );
        
        // 将客户端存储到会话映射中
        sessionMap.put(session.getId(), client);
        
        super.afterConnectionEstablished(session);
    }
    
    /**
     * 处理收到的文本消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        log.info("收到WebSocket消息，会话ID: {}, 消息: {}", session.getId(), message.getPayload());
        
        // 获取对应的讯飞客户端
        Ultra.XunfeiWebSocketClientStream client = sessionMap.get(session.getId());
        if (client != null) {
            // 将消息转发给讯飞API
            client.sendText(message.getPayload());
        } else {
            session.sendMessage(new TextMessage("error: 未找到对应的讯飞客户端"));
        }
    }
    
    /**
     * 连接关闭时的处理
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("WebSocket连接已关闭，会话ID: {}, 状态: {}", session.getId(), status);
        
        // 移除会话映射
        Ultra.XunfeiWebSocketClientStream client = sessionMap.remove(session.getId());
        if (client != null) {
            // 关闭讯飞客户端连接
            client.closeBlocking();
        }
        
        super.afterConnectionClosed(session, status);
    }
    
    /**
     * 处理传输错误
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误，会话ID: {}", session.getId(), exception);
        
        // 移除会话映射
        Ultra.XunfeiWebSocketClientStream client = sessionMap.remove(session.getId());
        if (client != null) {
            // 关闭讯飞客户端连接
            client.closeBlocking();
        }
        
        session.close(CloseStatus.SERVER_ERROR);
        super.handleTransportError(session, exception);
    }
    
    /**
     * 自定义流式回调处理器，将讯飞API的响应通过WebSocket发送给前端
     */
    private static class AIWebSocketStreamCallback implements org.ruoyi.interview.config.Ultra.StreamCallback {
        private final WebSocketSession session;
        
        public AIWebSocketStreamCallback(WebSocketSession session) {
            this.session = session;
        }
        
        @Override
        public void onMessage(String message) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(message));
                }
            } catch (Exception e) {
                log.error("发送WebSocket消息失败", e);
            }
        }
        
        @Override
        public void onComplete(String finalResult) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage("[DONE]"));
                }
            } catch (Exception e) {
                log.error("发送完成消息失败", e);
            }
        }
        
        @Override
        public void onError(Exception e) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage("error: " + e.getMessage()));
                }
            } catch (Exception ex) {
                log.error("发送错误消息失败", ex);
            }
        }
    }
}